#include <iostream>
#include <vector>
#include <array>
#include "Core/numlib.h"
#include "Core/TensorExpr.h"
#include "FiniteDiff/LevelOp.h"
#include "FiniteDiff/Multigrid.h"
#include "Embedded/EmbeddedGrid.h"
#include "Embedded/EmLevelOp.h"
#include "Embedded/EmFuncFiller.h"
#include "Embedded/EmIntergrid.h"
#include "Embedded/PoissonOp.h"
#include "Embedded/BlockRelaxation.h"
#include "../example/common/Ellipsoid.h"
#include "../example/common/TestUtility.h"

#if DIM == 3

const int Dim = SpaceDim;
using rVec = Vec<Real, Dim>;
using iVec = Vec<int, Dim>;
template class Tensor<Real, Dim>;

//======================================================================

const char *bcTypes = "DDDDDDD";

template <int Dim>
class PoissonLevelOp : public MGLevelOp<Dim>
{
public:
  using EmGrid = EmbeddedGrid<Dim>;
  using BaseClass = MGLevelOp<Dim>;

  PoissonLevelOp(const EmGrid &aGrid, Real w4j, const char *aBcTypes, bool isBottom = false)
      : BaseClass(aGrid, w4j, aBcTypes), eff(aGrid), elv(aGrid)
  {
    pPoissonOp = new PoissonOp<Dim>(aGrid, aBcTypes);
    pbr = new BlockRelaxation<Dim>(aGrid, pPoissonOp->matOp, 0.5);
    if(isBottom)
      pBottomSolver = new PoissonBottomSolver<Dim>(*pPoissonOp);
    else
      pBottomSolver = nullptr;
  }

  virtual ~PoissonLevelOp() {
    delete pBottomSolver;
    delete pbr;
    delete pPoissonOp;
  }

  Real computeNorm(const Tensor<Real, Dim> &phi) const override
  {
    return elv.NRM(phi, CellCentered, 0);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, const Tensor<Real, Dim-1> *bcData) const override
  {
    BaseClass::fillGhosts(phi, rhs, bcData);
    eff.transferBdryValues(phi, bcData[2*Dim]);
    pbr->apply(phi, rhs);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, int) const override
  {
    BaseClass::fillGhosts(phi, rhs, 0);
    eff.transferBdryValues(phi, 0);
    pbr->apply(phi, rhs);
  }

  void applySmoother(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &smoothed, const Tensor<Real, Dim> &rhs) const override
  {
    elv.RELAX_4TH(phi, rhs, smoothed, weightForJacobi);
  }

  void computeResidual(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &rsd, const Tensor<Real, Dim> &rhs) const override
  {
    pPoissonOp->apply(phi, rhs, rsd);
  }

  void solveBottom(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, int numBottomRelax) const override
  {
    assert(pBottomSolver != nullptr);
    pBottomSolver->solve(phi, rhs, 0);
  }

  const PoissonOp<Dim> &getPoissonOp() const { return *pPoissonOp; }

  using MGLevelOp<Dim>::getDomain;

protected:
  EmFuncFiller<Dim> eff;
  EmLevelOp<Dim> elv;
  PoissonOp<Dim> *pPoissonOp;
  PoissonBottomSolver<Dim> *pBottomSolver;
  BlockRelaxation<Dim> *pbr;
  using MGLevelOp<Dim>::weightForJacobi;
};

const PoissonLevelOp<Dim> *pDbgLevelOp;

template <class T0, class T1, class T2, int Dim>
std::array<Real,6> testNeumannPoisson(const std::vector<const MGLevelOp<Dim> *> &opHier,
                                      const std::vector<const Intergrid<Dim> *> &trHier,
                                      const T0 &PHI, const T1 &DPHI, const T2 &D2PHI)
{
  auto pLevelOp = dynamic_cast<const PoissonLevelOp<Dim> *>(opHier.front());
  const auto &poissonOp = pLevelOp->getPoissonOp();
  const auto &emGrid = poissonOp.emGrid;
  EmLevelOp<Dim> elv(emGrid);
  EmFuncFiller<Dim> eff(emGrid);
  const Box<Dim> gbx = emGrid.getGhostedBox();
  Tensor<Real, Dim> phi(gbx);
  Tensor<Real, Dim> err(gbx);
  Tensor<Real, Dim> rhs(gbx);
  Tensor<Real, Dim-1> cData[2*Dim+1];
  eff.fillPointValues(err, CellCentered, PHI);
  eff.fillPointValues(rhs, CellCentered, D2PHI);
  phi = 0.0;
  for(int d = 0; d < Dim; ++d)
    for (int m : {0, 1})
      cData[2 * d + m].resize(reduce(gbx, d));
  eff.fillBdryValues(cData, CellCentered, PHI);
  Vec<int,Dim-1> irBcSize(1);
  irBcSize[0] = emGrid.bdryNodes.size();
  cData[2*Dim].resize(irBcSize);
  eff.fillBdryValues(cData[2*Dim], "curvilinear", PHI); // Dirichlet boundary
//  eff.fillBdryGrads(cData[2*Dim], "curvilinear", DPHI); // Neumann boundary

  std::array<Real, 6> errnorm;

  // test the solution error
  Multigrid<Dim> mg(opHier, trHier);
  mg.setParam({3, 3, 0}, 40, 0.0, 1e-12);
  CPUTimer timer;
  timer.reset();
  mg.solve(phi, rhs, cData);
  std::cout << "Multigrid : " << std::defaultfloat << std::setprecision(5)
            << timer() << "s. " << std::endl;
  elv.zeroExtNodes(phi);
  err = err - phi;
  for(int p = 0; p <= 2; ++p)
    errnorm[p] = elv.NRM(err, CellCentered, p);

  // test the truncation error
  eff.fillPointValues(err, CellCentered, PHI);
  eff.transferBdryValues(err, cData[2*Dim]);
  poissonOp.apply(err, phi);
  err = rhs - phi;
  elv.zeroExtNodes(err);
  for(int p = 0; p <= 2; ++p)
    errnorm[3+p] = elv.NRM(err, CellCentered, p);
#ifndef NDEBUG
  if(get_dbglevel() >= 2) {
    err.dump(tmpos);
    tmpos.flush();
  }
#endif

  return errnorm;
}

template <class T0, class T1, class T2>
void doTest(const std::vector<int> &gridSize, int numCompHier, const T0 &PHI, const T1 &DPHI, const T2 &D2PHI)
{
  const Real tol = 1e-12;
  const int numGrid = gridSize.size();
  std::vector<const EmbeddedGrid<Dim> *> gridHier;
  std::vector<const PoissonLevelOp<Dim> *> opHier;
  std::vector<const EmIntergrid<Dim> *> trHier;
  rVec translate(0.5);
  Ellipsoid<Dim> geoObj(0.0, {0.25, 0.125, 0.25});
  for(int n = 0; n < numGrid; ++n) {
    const int N = gridSize[n];
    std::cout << "\n\nInitializing h = " << oneOver(N) << std::endl;
    RectDomain<Dim> rd(Box<Dim>(0, N-1), 1.0/N, CellCentered, 2);
    auto *pGrid = new EmbeddedGrid<Dim>(
        rd,
        [&](rVec q, rVec &closest, rVec &normal, bool &contained, Real tol) {
          q = q - translate;
          geoObj.check(q, closest, normal, contained, tol);
          closest = closest + translate;
          normal = -normal;
          contained = !contained;
        }, true, 0.1, tol);
#ifndef NDEBUG
    if(get_dbglevel() >= 2)
      pGrid->dump(tmpos);
#endif
    gridHier.push_back(pGrid);
    push_dbglevel(1);
    auto *pPoissonOp = new PoissonLevelOp<Dim>(*pGrid, 0.6, bcTypes, n == numGrid-1);
    std::cout.flush();
    pop_dbglevel();
    opHier.push_back(pPoissonOp);
    if(n != 0) {
      auto *pT = new EmIntergrid<Dim>(*pGrid, *(gridHier[n-1]));
      trHier.push_back(pT);
    }
  }

  std::vector<std::array<Real, 6>> errnorm(numCompHier);
  for(int n = numCompHier-1; n >= 0; --n) {
    std::vector<const MGLevelOp<Dim> *> _opHier(opHier.cbegin() + n, opHier.cend());
    std::vector<const Intergrid<Dim> *> _trHier(trHier.cbegin() + n, trHier.cend());
    std::cout << "\n\nTesting h = " << oneOver(gridSize[n]) << std::endl;
    errnorm[n] = testNeumannPoisson(_opHier, _trHier, PHI, DPHI, D2PHI);
  }
//
  for(int n = numGrid-1; n >= 0; --n) {
    if(n != 0)
      delete trHier[n-1];
    delete opHier[n];
    delete gridHier[n];
  }
  printConvergenceTable(&gridSize[0], errnorm);
}

int main()
{
  reset_dbglevel(1);
  std::cout << "SpaceDim = " << SpaceDim << std::endl;
  std::cout << "============================================================================\n";
  // the analytic solutions in Cartesian coordinates
  auto PHI = [&](rVec x) {
    return sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]) * sin(2*M_PI*x[2]);
  };
  auto DPHI = [&](rVec x) {
    return rVec { cos(2*M_PI*x[0]) * cos(2*M_PI*x[1]) * sin(2*M_PI*x[2]) * (2*M_PI),
                  -sin(2*M_PI*x[0]) * sin(2*M_PI*x[1]) * sin(2*M_PI*x[2]) * (2*M_PI),
                  sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]) * cos(2*M_PI*x[2]) * (2*M_PI) };
  };
  auto D2PHI = [&](rVec x) {
    return (-12*M_PI*M_PI) * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]) * sin(2*M_PI*x[2]);
  };
  //
  doTest({/*256, */128, 64, 32}, 3, PHI, DPHI, D2PHI);
  std::cout << std::endl;
  return 0;
}

#else

int main()
{
  std::cout << "============================================================================\n";
  std::cout << "testPoisson4 is skipped for 2D. " << std::endl;
  return 0;
}

#endif // DIM == 3
